#include <TCanvas.h>
#include <TH1.h>
#include <TRandom.h>
#include <TFile.h>
#include <TTree.h>
#include <iostream>
#include <iomanip>
#include <fstream>
#include <string>

// 获取导出目录
std::string GetExportPath()
{
    std::string path = __FILE__;
    path = path.substr(0, path.find("RootToPDF.C"));
    return path + "export/";
}

// 将ROOT直方图归一化为概率密度函数
void NormalizeHistToPDF(TH1 *hist)
{
    Double_t totalIntegral = hist->Integral();
    hist->Scale(1.0 / totalIntegral);
}

// 输出概率密度函数到文件
void OutputPDFToFile(TH1 *hist, const std::string &filename, const std::string &xTitle)
{
    std::ofstream outFile(filename);
    outFile << "pdf: tabular" << std::endl
            << "\t" << "x / " << xTitle << "\t" << "pdf" << std::endl;
    for (Int_t i = 1; i <= hist->GetNbinsX(); ++i)
    {

        if (hist->GetBinContent(i))
        {
            // 如果文件名包含"psf_number_distribution.properties"
            if (filename.find("psf_number_distribution.properties") != std::string::npos)
                outFile << std::fixed << "\t" << (int)hist->GetBinCenter(i) << "\t" << hist->GetBinContent(i) << std::endl;
            else
                outFile << std::fixed << std::setprecision(18) << "\t" << hist->GetBinCenter(i) << "\t" << hist->GetBinContent(i) << std::endl;
        }
    }
}

// 绘制并保存直方图
void DrawAndSaveHist(TH1 *hist, const std::string &xTitle, const std::string &filename, bool drawPlots, bool savePlots)
{
    if (drawPlots || savePlots)
    {
        TCanvas *canvas = new TCanvas();
        hist->GetXaxis()->SetTitle(xTitle.c_str());
        hist->GetXaxis()->SetTitleOffset(1.4);
        hist->GetYaxis()->SetTitle("pdf");
        hist->GetYaxis()->SetTitleOffset(1.4);
        hist->GetZaxis()->SetTitleOffset(1);
        hist->Draw();
        if (savePlots)
            canvas->SaveAs(filename.c_str());
        // delete canvas;
    }
}

// 从ROOT文件中读取直方图数据并输出为概率密度函数
void RootToPDF()
{
    std::vector<TH1 *> histos = {
        new TH1D("Energy", "", 50, 2.0, 3.47), // 光子能量
        new TH1D("Time", "", 100, 0, 50),      // 光子时间
        new TH1D("PosX", "", 100, -0.5, 0.5),  // 光子位置
        new TH1D("PosY", "", 100, -0.5, 0.5),  // 光子位置
        new TH1D("DirX", "", 100, -1, 1),      // 光子方向
        new TH1D("DirY", "", 100, -1, 1),      // 光子方向
        new TH1D("DirZ", "", 100, 0, 1),       // 光子方向
        new TH1D("DirPhi", "", 180, 0, 180),   // 光子方向
        new TH1D("DirTheta", "", 180, 0, 180), // 光子方向
        new TH1I("NPhoton", "", 1000, 0, 1000) // 光子数
    };
    TFile *f = new TFile("psfhisto.root", "open");
    TTree *tree = (TTree *)f->Get("ntuple");
    int entries = int(tree->GetEntries());
    std::cout << "Entries:" << entries << std::endl;

    std::vector<double> *fVecEnergy = nullptr;
    std::vector<double> *fVecTime = nullptr;
    std::vector<double> *fVecPosX = nullptr;
    std::vector<double> *fVecPosY = nullptr;
    std::vector<double> *fVecDirX = nullptr;
    std::vector<double> *fVecDirY = nullptr;
    std::vector<double> *fVecDirZ = nullptr;
    std::vector<int> *fVecNPhoton = nullptr;
    std::vector<double> *fVecDirPhi = nullptr;
    std::vector<double> *fVecDirTheta = nullptr;

    tree->SetBranchAddress("Energy", &fVecEnergy);
    tree->SetBranchAddress("Time", &fVecTime);
    tree->SetBranchAddress("PosX", &fVecPosX);
    tree->SetBranchAddress("PosY", &fVecPosY);
    tree->SetBranchAddress("DirX", &fVecDirX);
    tree->SetBranchAddress("DirY", &fVecDirY);
    tree->SetBranchAddress("DirZ", &fVecDirZ);
    tree->SetBranchAddress("NPhoton", &fVecNPhoton);
    tree->SetBranchAddress("DirPhi", &fVecDirPhi);
    tree->SetBranchAddress("DirTheta", &fVecDirTheta);

    // 遍历所有条目，填充直方图
    for (size_t i = 0; i < entries; i++)
    {
        tree->GetEntry(i);
        for (size_t j = 0; j < fVecEnergy->size(); j++)
        {
            histos[0]->Fill(fVecEnergy->at(j));
            histos[1]->Fill(fVecTime->at(j));
            histos[2]->Fill(fVecPosX->at(j));
            histos[3]->Fill(fVecPosY->at(j));
            histos[4]->Fill(fVecDirX->at(j));
            histos[5]->Fill(fVecDirY->at(j));
            histos[6]->Fill(fVecDirZ->at(j));
            histos[7]->Fill(fVecDirTheta->at(j));
            histos[8]->Fill(fVecDirPhi->at(j));
        }
        histos[9]->Fill(fVecNPhoton->at(0));
    }

    // 打印光子数的均值
    std::cout << "NPhotonMean:" << histos[9]->GetMean() << std::endl;

    // 归一化为概率密度函数
    for (size_t i = 0; i < 9; i++)
    {
        NormalizeHistToPDF(histos[i]);
    }
    std::string exportpath = GetExportPath();
    // 输出概率密度函数到文件
    OutputPDFToFile(histos[0], exportpath + "prop/psf_energy_distribution.properties", "eV");
    OutputPDFToFile(histos[1], exportpath + "prop/psf_time_distribution.properties", "ns");
    OutputPDFToFile(histos[7], exportpath + "prop/psf_phi_distribution.properties", "deg");
    OutputPDFToFile(histos[8], exportpath + "prop/psf_theta_distribution.properties", "deg");
    OutputPDFToFile(histos[9], exportpath + "prop/psf_number_distribution.properties", "");

    // 重新绘制归一化后的直方图并保存为svg
    bool drawPlots = true; // 控制是否绘制图像
    bool savePlots = true; // 控制是否保存图像

    DrawAndSaveHist(histos[0], "Energy/eV", exportpath + "pic/psf_energy_distribution.svg", drawPlots, savePlots);
    DrawAndSaveHist(histos[1], "Time/ns", exportpath + "pic/psf_time_distribution.svg", drawPlots, savePlots);
    DrawAndSaveHist(histos[2], "PosX/mm", exportpath + "pic/psf_positionx_distribution.svg", drawPlots, savePlots);
    DrawAndSaveHist(histos[3], "PosY/mm", exportpath + "pic/psf_positiony_distribution.svg", drawPlots, savePlots);
    DrawAndSaveHist(histos[4], "DirX", exportpath + "pic/psf_directionx_distribution.svg", drawPlots, savePlots);
    DrawAndSaveHist(histos[5], "DirY", exportpath + "pic/psf_directiony_distribution.svg", drawPlots, savePlots);
    DrawAndSaveHist(histos[6], "DirZ", exportpath + "pic/psf_directionz_distribution.svg", drawPlots, savePlots);
    DrawAndSaveHist(histos[7], "DirPhi", exportpath + "pic/psf_phi_distribution.svg", drawPlots, savePlots);
    DrawAndSaveHist(histos[8], "DirTheta", exportpath + "pic/psf_theta_distribution.svg", drawPlots, savePlots);
    DrawAndSaveHist(histos[9], "NPhoton", exportpath + "pic/psf_number_distribution.svg", drawPlots, savePlots);
}
